import os
import time
import random
import functools
import itertools
import argparse
import numpy as np


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', type=str, default='/tcdata')
    parser.add_argument('--output_dir', type=str, default='cls_labels')
    parser.add_argument('--mode', type=int, default=0)
    args = parser.parse_args()
    return args


def split_image_data(fname, output_dir, ratio=0.9):
    random.seed(100)
    with open(fname, 'r', encoding='utf8') as fh:
        lines = fh.readlines()

    num = len(lines)
    num_train = int(ratio * num)
    indices = list(range(num))
    random.shuffle(indices)
    flag = np.zeros(num, dtype=bool)
    flag[indices[:num_train]] = True

    basename = os.path.splitext(os.path.basename(fname))[0]
    os.makedirs(output_dir, exist_ok=True)
    train_fname = os.path.join(output_dir, basename + '_train.txt')
    train_file = open(train_fname, 'w', encoding='utf8')
    val_fname = os.path.join(output_dir, basename + '_val.txt')
    val_file = open(val_fname, 'w', encoding='utf8')

    for i in range(num):
        if flag[i]:
            train_file.write(lines[i])
        else:
            val_file.write(lines[i])
    train_file.close()
    val_file.close()


def main():
    args = parse_args()
    data_dir = args.data_dir
    output_dir = args.output_dir
    data_sets = [d for d in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, d))]
    os.makedirs(output_dir, exist_ok=True)
    for set_name in data_sets:
        files = os.listdir(os.path.join(data_dir, set_name))
        outname = os.path.join(output_dir, set_name+'.txt')
        with open(outname, 'w', encoding='utf-8') as fh:
            for name in files:
                fh.write('{}/{}\t0\n'.format(set_name, name))
        split_image_data(outname, output_dir)



if __name__ == '__main__':
    main()
